"""
用于将来自论文[1] Yang H. , Xu X. , Ma Y. , et al.CraterDANet: A Convolutional Neural Network for Small-Scale Crater Detection via Synthetic-to-Real Domain Adaptation[J/OL].IEEE Trans. Geosci. Remote Sensing,2022,60:1-12

的数据集处理为需要的子块格式，其数据下载源为：https://github.com/yizuifangxiuyh/Lunar_Crater_Detection_Data
"""

import os
import numpy as np
import cv2


def div_txt(row, col, bbox, base_batch):
    """
    将某图对应的bbox标注文件划分到子图上
    """
    inliers = []
    for b in bbox:
        x, y, w, h = b
        if (
            x >= col
            and y >= row
            and x + w <= col + base_batch
            and y + h <= row + base_batch
        ):
            inliers.append([x - col, y - row, w, h])
    return inliers


dataset_dir = "/home/a804_cbf/Code/Lunar_Crater_Detection_Data/LRO_DATA/test"
output_dir = "/disk527/sdb1/a804_cbf/datasets/CraterDANet/test"

if not os.path.exists(output_dir):
    os.makedirs(output_dir)

for file in os.listdir(dataset_dir):
    if file.endswith(".png"):
        image = cv2.imread(os.path.join(dataset_dir, file), cv2.IMREAD_GRAYSCALE)
        with open(os.path.join(dataset_dir, file.replace(".png", ".txt"))) as f:
            bbox = []
            f.readline()
            for line in f:
                line = line.strip().split(",")
                bbox.append([round(float(x)) for x in line])
        # 存储以xywh格式的bbox
        bbox = np.array(bbox)[:, 1:5]
        # 分为四分之一块，每块大小为200*200
        # 以滑窗法提取patch，重叠区域大小为180，即每次滑动20
        base_batch = 200
        stride = 20
        for i in range(0, image.shape[0] - base_batch, stride):
            for j in range(0, image.shape[1] - base_batch, stride):
                patch = image[i : i + base_batch, j : j + base_batch]
                cv2.imwrite(os.path.join(output_dir, f"{i}_{j}_{file}"), patch)
                sub_bbox = div_txt(i, j, bbox, base_batch)
                with open(
                    os.path.join(output_dir, f"{i}_{j}_{file.replace('.png', '.txt')}"),
                    "w",
                ) as f:
                    f.write("x,y,w,h\n")
                    for b in sub_bbox:
                        f.write(f"{b[0]},{b[1]},{b[2]},{b[3]}\n")
